#----------------------------------------------------------
# File SolfecPanel.py
# 
# The following Blender panel is drawn in the the Scene tab
# of the Properties Window, and is meant to help the user
# export a scene into a Solfec input file.
#
# Jedediyah Williams
#      


#----------------------------------------------------------
# All the python code for the Blender UI and Solfec 
# translation is included in this single file.
# 
# Important sections of code:
#   1. def initSceneProperties(scn)
#       Creates the variables for storing Solfec properties.
#       These should all be in bpy.types.Scene.[variable_name]
#   2. class UIPanel(bpy.types.Panel):
#       Defines the layout of the panel 
#   3. class OBJECT_OT_Export_to_Solfec(bpy.types.Operator):
#       Handles user button click "Export scene to Solfec."
# 

import os
import math
import mathutils
import bpy
from bpy.props import *
 
#----------------------------------------------------------
#    Store properties in the active scene.
#
def initSceneProperties(scn):
 
    ### Project Properties ###
    
    # Project name
    bpy.types.Scene.Solfec_projectname = StringProperty(
        name = "Project Title",
        description = "Name of Solfec Project")
    scn['Solfec_projectname'] = "Solfec_Project"
    
    # Directory path to place Solfec file
    bpy.types.Scene.Solfec_dirpath = StringProperty(
        name="Directory",
        description = "The directory to place the Solfec input file",
        subtype = "DIR_PATH")
        
    # File overwrite preference
    bpy.types.Scene.Solfec_overwriteFile = BoolProperty(
        name = "Overwrite existing files.",
        description = "If true, will overwrite file if it exists.",
        default = False)
    
    # Step size
    bpy.types.Scene.Solfec_stepsize = FloatProperty(
        name = "Step size", 
        description = "Time step size during Solfec simulation.",
        default = 0.001,
        min = 1E-10,
        max = 100)
    scn['Solfec_stepsize'] = 0.001
 
    # Stop time
    bpy.types.Scene.Solfec_stoptime = FloatProperty(
        name = "Stop time",
        description = "When to stop Solfec simulation.",
        default = 10.0,
        min = 0.0,
        max = 1E10000)
    scn['Solfec_stoptime'] = 10.0
 
    # Solver
    bpy.types.Scene.Solfec_solver = EnumProperty(
        items = [('Gauss-Seidel', 'Gauss-Seidel (recommended)', 'Gauss-Seidel'), 
                 ('Penalty', 'Penalty', 'Penalty'),
                 ('quasi-Newton', 'quasi-Newton', 'quasi-Newton')],
        name = "Solver")
    scn['Solfec_solver'] = 0
    
    
    ### MATERIAL properties ###
     
    # Friction
    bpy.types.Scene.Solfec_friction = FloatProperty(
        name = "Friction",
        description = "[Surface Material] Friction for all Solfec bodies.",
        default = 0.5,
        min = 0.0,
        max = 10.0)
    scn['Solfec_friction'] = 0.5
    
    # Restitution
    bpy.types.Scene.Solfec_restitution = FloatProperty(
        name = "Restitution",
        description = "[Surface Material] Velocity restitution for all Solfec bodies.",
        default = 0.0,
        min = 0.0,
        max = 1.0)
    scn['Solfec_restitution'] = 0.0
    
    # Young's modulus
    bpy.types.Scene.Solfec_young = FloatProperty(
        name = "Young's Modulus",
        description = "[Bulk Material] Young's modulus.",
        default = 1E6,
        min = 0.0,
        max = 2E8)
    scn['Solfec_young'] = 1E6
       
    # Poisson's coefficient
    bpy.types.Scene.Solfec_poisson = FloatProperty(
        name = "Poisson's Coefficient",
        description = "[Bulk Material] Poisson's coefficient.",
        default = 0.25,
        min = 0.0,
        max = 0.6)
    scn['Solfec_poisson'] = 0.25
    
    # Density
    bpy.types.Scene.Solfec_density = FloatProperty(
        name = "Density",
        description = "[Bulk Material] Density.",
        default = 1E3,
        min = 0.0,
        max = 1E9)  # Ok, so that's the density of a white dwarf.  Excessive?
    scn['Solfec_density'] = 1E3
    

    # Finished defining Blender-Solfec data
    return

initSceneProperties(bpy.context.scene)
 
#----------------------------------------------------------
#    Panel in PROPERTIES region
#
class UIPanel(bpy.types.Panel):
    bl_label = "Solfec Export"
    bl_space_type = "PROPERTIES"
    bl_region_type = "WINDOW"
    bl_context = "physics"  # Simply more convenient to have it in Physics
 
    def draw(self, context):
        layout = self.layout
        scn = context.scene
        
        layout.label(text = "Project Properties:", icon = "GROUP")
        box = layout.box()
        box.prop(scn, 'Solfec_projectname')
        row = box.row()
        row.prop(scn, 'Solfec_dirpath')
        row = box.row()
        row.alignment = 'RIGHT'
        row.prop(scn, 'Solfec_overwriteFile')
        row = box.row()
        row.prop(scn, 'Solfec_stepsize')
        row.prop(scn, 'Solfec_stoptime')
        row = box.row()
        row.prop(scn, 'Solfec_solver')
        row = layout.row()
        
        row.label(text = "Material Properties:", icon = "MOD_UVPROJECT")
        row = layout.row()
        box = row.box()
        box.prop(scn, 'Solfec_friction')
        box.prop(scn, 'Solfec_restitution')
        box.prop(scn, 'Solfec_young')
        box.prop(scn, 'Solfec_poisson')
        box.prop(scn, 'Solfec_density')
        
        layout.label("Create a Solfec input file:", icon = "TEXT")
        #row = layout.row()
        #row.label(" ", icon = "NONE")
        box = layout.box()
        row = box.row()
        row.operator("export.solfec")
        row.label("", icon = "FILE_TICK")
 
#----------------------------------------------------------
#    The "Export" button starts the export process.
#
class OBJECT_OT_Export_to_Solfec(bpy.types.Operator):
    bl_idname = "export.solfec"
    bl_label = "Export scene to Solfec"
 
    def execute(self, context):
        scn = context.scene
        print()
        
        # Output to Blender console
        print("Generating Solfec file from blender...")
        projName = scn.Solfec_projectname
        filepath = scn.Solfec_dirpath + projName + ".py"
        print(filepath)
        
        # Check if file already exists
        if os.path.exists(filepath) and not scn.Solfec_overwriteFile:
            self.report('ERROR', scn.Solfec_projectname + ".py already exists.  Solfec export canceled.")
            return{'FINISHED'}
            
        
        f = open(filepath, 'w')
        # Create Solfec file #############################################
        f.write("# A Solfec file generated in Blender.\n\n")
        f.write((""))
        f.write("step = " + str(round(scn.Solfec_stepsize,9)) + "\n")
        f.write("stop = " + str(scn.Solfec_stoptime) + "\n")

        f.write("output = 'out/" + projName + "'\n");
        f.write("solfec = SOLFEC('DYNAMIC', step, output)\n");
        f.write("bulkmat = BULK_MATERIAL (solfec, model='KIRCHHOFF'" + \
            ", young = " + str(scn.Solfec_young) + \
            ", poisson = " + str(scn.Solfec_poisson) + \
            ", density = " + str(scn.Solfec_density) + ")\n");
        
        # Surface Material
        if scn.Solfec_solver == "Penalty":
            f.write("surfmat = SURFACE_MATERIAL (solfec, model = 'SPRING_DASHPOT'" + \
            ", friction = " + str(scn.Solfec_friction) + \
            ", restitution = " + str(scn.Solfec_restitution) + \
            "spring = 1E+9, dashpot = 1E6)\n");         # Should be user input.
        
        else:    
            f.write("surfmat = SURFACE_MATERIAL (solfec, model = 'SIGNORINI_COULOMB'" + \
                ", friction = " + str(scn.Solfec_friction) + \
                ", restitution = " + str(scn.Solfec_restitution) + ")\n");
            
        # Body attributes
        b_vid = 1;
        b_sid = 1;
        
        # Generate Solfec formatted bodies 
        for bodyObject in bpy.data.objects.values():
            if bodyObject.type == "MESH":
                makeBody(f, bodyObject, b_vid, b_sid)
                b_vid += 1;     # Don't know if this works.  
                b_sid += 1;
                
        f.write("\n# CONSTRAINTS.\n")
        # Generate Solfec formatted constraints
        for bodyObject in bpy.data.objects.values():
            if bodyObject.type == "MESH":
                for const in bodyObject.constraints.values():
                    makeConstraint(f, bodyObject, const)
                    
        f.write("\n# Define constraints based on Blender parent-child objects. \n")
        # Parent-child constraints.  In blender, child objects are rigidly constrained to thier parent.  
        for bodyObject in bpy.data.objects.values():
            for chld in bodyObject.children:
                f.write("# Parent: " + SolfecSafeName(bodyObject.name) + \
                         ", Child: " + SolfecSafeName(chld.name) + "\n")
                makeChildConstraint(f, bodyObject, chld)
                f.write("\n")
                
        
        f.write("\n# Define gravity and run.\n");
        f.write("GRAVITY(solfec, (0,0,-9.81))\n");
        
        # Output
        f.write("OUTPUT(solfec, 0.01)\n")   # Avoid writing too much to disk
        
        # Solver
        if scn.Solfec_solver == "Penalty":
            f.write("ps = PENALTY_SOLVER() \n")
            f.write("RUN(solfec, ps, stop) \n")
        elif scn.Solfec_solver == "quasi-Newton":
            f.write("ns = NEWTON_SOLVER() \n")
            f.write("RUN(solfec, ns, stop) \n")
        else:   
            f.write("gs = GAUSS_SEIDEL_SOLVER(1E-4,200) \n")
            f.write("RUN(solfec, gs, stop) \n");
        
        
        # Done with Solfec file #############################################
        f.close();
        
        print("done."); 
        
        
        return{'FINISHED'}    



#---------------------------------------------------------- 
# Transform vector V to world frame, given transformation matrix T
def vertWorld(T, V):
    # Initialize the world frame vector Vw
    Vw = mathutils.Vector((0.0, 0.0, 0.0))
    # Multiply R*V
    Vw[0] = T[0][0]*V[0] + T[0][1]*V[1] + T[0][2]*V[2]
    Vw[1] = T[1][0]*V[0] + T[1][1]*V[1] + T[1][2]*V[2]
    Vw[2] = T[2][0]*V[0] + T[2][1]*V[1] + T[2][2]*V[2]
    # Add P
    Vw[0] += T[0][3]
    Vw[1] += T[1][3]
    Vw[2] += T[2][3]
    # Return Vw
    return Vw

#---------------------------------------------------------- 
# A utility function
def SolfecSafeName(name):
    return name.replace(".","").replace("-","_");

#---------------------------------------------------------- 
# Generate Solfec formatted body
def makeBody(f, bodyObject, b_vid, b_sid):
    # Get mesh object
    body = bodyObject.data
    T = bodyObject.matrix_world.transposed()
    # Create solfec model of Blender mesh
    bodyName = SolfecSafeName(bodyObject.name)
    print(" Generating model of " + str(body) + "...")
    output = "\n#" + str(body) + "\n";
    
    # Case of collision type SPHERE
    if bodyObject.game.collision_bounds_type == 'SPHERE':
        loc = bodyObject.location
        output += bodyName + " = SPHERE((" + str(loc[0]) + ", " + str(loc[1]) +", " + str(loc[2]) + "), "
        output += str(bodyObject.game.radius * bodyObject.scale[0]) + ", " + str(b_vid) + ", " + str(b_sid) + ")\n"
        f.write(output)
        
    # Case of collision type Triangle Mesh
    elif bodyObject.game.collision_bounds_type == 'TRIANGLE_MESH':
        output += bodyName + " = MESH ([\n"
        f.write(output)
        for i in range(len(body.vertices)-1):
            v = body.vertices[i].co
            v = vertWorld(T,v)
            output = str(v[0]) + ", " + str(v[1]) + ", " + str(v[2]) + ",\n"
            f.write(output)
        v = body.vertices[len(body.vertices)-1].co
        v = vertWorld(T,v)
        output = str(v[0]) + ", " + str(v[1]) + ", " + str(v[2]) + "],\n"
        f.write(output)
        
        # Mesh should be either a tetrahedron, pyramid, wedge, or hexahedron.
        nv = len(body.vertices)
        ne = len(body.edges)
        nf = len(body.faces)
        
        if nv == 8 and ne == 12 and nf == 6:  # Hexahedron
            output = "[8, 0, 1, 2, 3, 4, 5, 6, 7, 4], " + str(b_vid) + ")\n"
        elif nv == 4 and ne == 6 and nf == 4: # Tetrahedron
            output = "[4, 0, 1, 2, 3, 1], " + str(b_vid) + ")\n"
        elif nv == 5 and ne == 8 and nf == 5: # Pyramid
            output = "[5, 0, 1, 2, 3, 4, 2], " + str(b_vid) + ")\n"
        elif nv == 6 and ne == 9 and nf == 5: # Wedge
            output = "[6, 0, 1, 2, 3, 4, 5, 3], " + str(b_vid) + ")" 
        else:
            self.report('ERROR', "Attempted to apply a triangle mesh to an object other than Tetrahedron, Pyramid, Hexahedron, or Wedge")
        f.write(output)
        
        # TETRAHEDRALIZE
        # TODO - allow user to specify attributes...
        #f.write("TETRAHEDRALIZE(" + SolfecSafeName(bodyName) + ", 'out/" + bpy.context.scene.Solfec_projectname + "/meshes/" + bodyName + ".dat',.02,1.5)\n")
        
    # Case of collision type HULL
    else:
        output += bodyName + " = HULL(["
        f.write(output)
        # For n-1 vertices 
        for i in range(len(body.vertices.values())-1):
            v = body.vertices.values()[i].co
            v = vertWorld(T,v)
            output = "\n "
            output = output + str(v[0]) + ", " + str(v[1]) + ", " + str(v[2]) + ","
            f.write(output)
        # Take care of the last (nth) vertex
        v = body.vertices.values()[len(body.vertices.values())-1].co
        v = vertWorld(T,v)
        output = "\n "
        output = output + str(v[0]) + ", " + str(v[1]) + ", " + str(v[2]) + ","
        output += "], " + str(b_vid) + ", " + str(b_sid) + ")\n"
        f.write(output)
    
    # Determine physics type i.e. Static or Rigid body
    if bodyObject.game.physics_type == 'STATIC':
        bodyPhysicsType = 'OBSTACLE'
    elif bodyObject.game.physics_type == 'SOFT_BODY':
        bodyPhysicsType = 'FINITE_ELEMENT' #'PSEUDO_RIGID'
    else:
        bodyPhysicsType = 'RIGID'
    output = bodyName + "_body = BODY (solfec, '" + bodyPhysicsType + "', " + bodyName + ", bulkmat)\n\n";
    f.write(output);
    
    ############################## A TEST
    f.write(bodyName + "_body.nodecontact='ON' \n")

#---------------------------------------------------------- 
# Generate Solfec constraints, given file f, body b, and constraint c
def makeConstraint(f,b,c):
     
    # Target body
    t = c.target
    
    # HINGE joint constraint
    if c.pivot_type == "HINGE":
        # Rotation matrices
        R0j = b.matrix_world.transposed().to_3x3()
        Rx = mathutils.Matrix.Rotation(-c.axis_x,3,'X')
        Ry = mathutils.Matrix.Rotation(-c.axis_y,3,'Y')
        Rz = mathutils.Matrix.Rotation(-c.axis_z,3,'Z')

        # Get direction of pivot's x-axis (axis of rotation in Blender)        
        X = mathutils.Vector((1.0,0,0,0))   
        Z = mathutils.Vector((0,0,1.0))
        
        Vx = mathutils.Vector((0,0,0))
        Vz = mathutils.Vector((0,0,0))
        
        R = Rx*Ry*Rz*R0j
        Vx[0] = R[0][0]*X[0] + R[0][1]*X[1] + R[0][2]*X[2]
        Vx[1] = R[1][0]*X[0] + R[1][1]*X[1] + R[1][2]*X[2]
        Vx[2] = R[2][0]*X[0] + R[2][1]*X[1] + R[2][2]*X[2]
        
        Vz[0] = R[0][0]*Z[0] + R[0][1]*Z[1] + R[0][2]*Z[2]
        Vz[1] = R[1][0]*Z[0] + R[1][1]*Z[1] + R[1][2]*Z[2]
        Vz[2] = R[2][0]*Z[0] + R[2][1]*Z[1] + R[2][2]*Z[2]
        
        # Find points for Solfec constraints
        D = .6  # Distance between ith body's constraint points 
        p0i = b.location
        p0j = t.location
        
        pi1 = p0i + (D/2)*Vx
        pi2 = p0i - (D/2)*Vx
        pj1 = p0j + (D/2)*Vx
        pj2 = p0j - (D/2)*Vx
        
        # Simulate a hinge joint by
        output = "PUT_RIGID_LINK(" + SolfecSafeName(b.name) + "_body, " + SolfecSafeName(t.name) + "_body, (";
        output += str(pi1[0])+","+str(pi1[1])+","+str(pi1[2])+"), ("+str(pj1[0])+","+str(pj1[1])+","+str(pj1[2])+"))\n";
        f.write(output)
    
        output = "PUT_RIGID_LINK(" + SolfecSafeName(b.name) + "_body, " + SolfecSafeName(t.name) + "_body, (";
        output += str(pi2[0])+","+str(pi2[1])+","+str(pi2[2])+"), ("+str(pj2[0])+","+str(pj2[1])+","+str(pj2[2])+"))\n";
        f.write(output)
        
    # BALL joint constraint
    elif c.pivot_type == "BALL":
        p1 = b.location
        p2 = t.location
        output = "PUT_RIGID_LINK(" + SolfecSafeName(b.name) + "_body, " + SolfecSafeName(t.name) + "_body, (";
        output += str(p1[0])+","+str(p1[1])+","+str(p1[2])+"), ("+str(p2[0])+","+str(p2[1])+","+str(p2[2])+"))\n";
        f.write(output)
    
    else:
        print("WARNING: Constraint of type " + str(c) + "is not implemented.")
        
    
#----------------------------------------------------------
#    In Blender, child objects are rigidly constrained to 
#    their parent.  This method generates these constraints
#    in Solfec using 3 rigid spherical joints.       
def makeChildConstraint(f, prnt, chld):
    
    print("Creating parent-child constraint; Parent: " + prnt.name + ", Child: " + chld.name)
    
    # Initialize vector from parent to child
    V = mathutils.Vector((0,0,0))
    V[0] = chld.location[0] - prnt.location[0]
    V[1] = chld.location[1] - prnt.location[1]
    V[2] = chld.location[2] - prnt.location[2]
    
    Vz = V/math.pow(V[0]*V[0] + V[1]*V[1] + V[2]*V[2],0.5)
    
    # Get an arbitrary orthogonal vector to V
    Vx = mathutils.Vector((1,0,0))
    dp = Vz[0]*Vx[0] + Vz[1]*Vx[1] + Vz[2]*Vx[2]
    if dp > 1.414:  # Rotate around Y; (x, y, z) -> (z, y, -x)
        Vx[0] = Vz[2]
        Vx[1] = Vz[1]
        Vx[2] = -Vz[0]
    else:           # Rotate around X; (x, y, z) -> (x, z, -y)
        Vx[0] = Vz[0]
        Vx[1] = Vz[2]
        Vx[2] = -Vz[1]

    # Find Vy orthonormal to Vx, Vz
    Vy = mathutils.Vector((0,0,0))
    Vy[0] = Vz[1]*Vx[2] - Vz[2]*Vx[1]
    Vy[1] = Vz[2]*Vx[0] - Vz[0]*Vx[2]
    Vy[2] = Vz[0]*Vx[1] - Vz[1]*Vx[0]
            
    # Create 3 rigid constraints between parent and child.
    dy = .6     # Defines how far apart the constraints will be placed.
    x = dy*math.cos(math.pi/3)
    y = dy/2

    # Dither around COM of parent
    p1a = p1b = p1c = prnt.location
    p1a = p1a + Vy*y
    p1b = p1b - Vx*x -Vy*y
    p1c = p1c + Vx*x -Vy*y
    
    p2a = p1a + V
    p2b = p1b + V
    p2c = p1c + V
    
    parent_body = SolfecSafeName(prnt.name)
    child_body = SolfecSafeName(chld.name)
    
    f.write("PUT_RIGID_LINK(" + parent_body + "_body, " + child_body + "_body, (" + \
            str(p1a[0]) + ", " + str(p1a[1]) + ", " + str(p1a[2]) + "), (" + \
            str(p2a[0]) + ", " + str(p2a[1]) + ", " + str(p2a[2]) + ") ) \n");
            
    f.write("PUT_RIGID_LINK(" + parent_body + "_body, " + child_body + "_body, (" + \
            str(p1b[0]) + ", " + str(p1b[1]) + ", " + str(p1b[2]) + "), (" + \
            str(p2b[0]) + ", " + str(p2b[1]) + ", " + str(p2b[2]) + ") ) \n");
            
    f.write("PUT_RIGID_LINK(" + parent_body + "_body, " + child_body + "_body, (" + \
            str(p1c[0]) + ", " + str(p1c[1]) + ", " + str(p1c[2]) + "), (" + \
            str(p2c[0]) + ", " + str(p2c[1]) + ", " + str(p2c[2]) + ") ) \n");

 
#----------------------------------------------------------
#    Registration
bpy.utils.register_module(__name__)
